package com.happy3w.playground.thread;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.stream.Stream;

public class Manager<T> {
    private final int threadCount;

    public Manager(int threadCount) {
        this.threadCount = threadCount;
    }

    public void startWork(Stream<T> dataStream,
                          T endFlag,
                          Consumer<T> consumer) {
        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
        LinkedBlockingQueue<T> queue = new LinkedBlockingQueue<>(100);
        AtomicInteger workerCount = new AtomicInteger(0);

        try {
            startWorkers(consumer, endFlag, executor, queue, workerCount);
            dataStream.forEach(queue::add);
        } finally {
            executor.shutdown();
            try {
                fillEndFlag(queue, endFlag);
                while (!queue.isEmpty() || workerCount.get() > 0) {
                    Thread.sleep(200);
                }
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    private void fillEndFlag(LinkedBlockingQueue<T> queue, T endFlag) throws InterruptedException {
        for (int i = 0; i < threadCount; i++) {
            queue.put(endFlag);
        }
    }

    private void startWorkers(Consumer<T> consumer, T endFlag, ExecutorService executor, LinkedBlockingQueue<T> queue, AtomicInteger workerCount) {
        for (int i = 0; i < threadCount; i++) {
            executor.execute(new Worker(consumer, endFlag, queue, workerCount));
        }
    }
}
